"""
核反应堆子通道分析程序的绘图函数
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.animation import FuncAnimation
import os

plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文显示
plt.rcParams['axes.unicode_minus'] = False  # 设置负号显示

def plot_results(analysis):
    """
    绘制计算结果的可视化图
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        # 创建网格点
        x = np.arange(analysis.n_axial)
        y = np.arange(analysis.n_channels)
        X, Y = np.meshgrid(x, y)
        
        # 检查数据是否为空
        if np.all(analysis.void_fraction == 0) or np.all(analysis.pressure == 0):
            print("警告：检测到空数据，请确保计算已完成")
            return
            
        # 创建两个图形窗口
        fig1 = plt.figure(figsize=(20, 15))
        fig2 = plt.figure(figsize=(20, 15))
        
        # 第一个图形窗口：液相参数
        axes1 = []
        for i in range(6):
            ax = fig1.add_subplot(2, 3, i+1, projection='3d')
            axes1.append(ax)
            
        # 第二个图形窗口：气相参数
        axes2 = []
        for i in range(6):
            ax = fig2.add_subplot(2, 3, i+1, projection='3d')
            axes2.append(ax)
            
        # 绘制液相参数
        plots = [
            (analysis.alpha_l, '液相分数分布'),
            (analysis.u_l, '液相轴向速度分布 (m/s)'),
            (analysis.v_l, '液相横向速度分布 (m/s)'),
            (analysis.T_l, '液相温度分布 (K)'),
            (analysis.pressure, '压力分布 (Pa)'),
            (analysis.q_wl if hasattr(analysis, 'q_wl') else np.zeros_like(analysis.pressure), 
             '液相壁面热流密度 (W/m²)')
        ]
        
        for ax, (data, title) in zip(axes1, plots):
            surf = ax.plot_surface(X, Y, data.T, cmap=cm.viridis)
            ax.set_title(title)
            ax.set_xlabel('轴向位置')
            ax.set_ylabel('子通道编号')
            fig1.colorbar(surf, ax=ax, shrink=0.5, aspect=5)
            
        # 绘制气相参数
        plots = [
            (analysis.alpha_v, '气相分数分布'),
            (analysis.u_v, '气相轴向速度分布 (m/s)'),
            (analysis.v_v, '气相横向速度分布 (m/s)'),
            (analysis.T_v, '气相温度分布 (K)'),
            (analysis.gamma if hasattr(analysis, 'gamma') else np.zeros_like(analysis.pressure), 
             '相变率 (kg/m³·s)'),
            (analysis.q_wv if hasattr(analysis, 'q_wv') else np.zeros_like(analysis.pressure), 
             '气相壁面热流密度 (W/m²)')
        ]
        
        for ax, (data, title) in zip(axes2, plots):
            surf = ax.plot_surface(X, Y, data.T, cmap=cm.viridis)
            ax.set_title(title)
            ax.set_xlabel('轴向位置')
            ax.set_ylabel('子通道编号')
            fig2.colorbar(surf, ax=ax, shrink=0.5, aspect=5)
            
        # 设置视角和网格
        for ax in axes1 + axes2:
            ax.view_init(elev=20, azim=45)
            ax.grid(True)
            
        plt.tight_layout()
        plt.show()
        
    except Exception as e:
        print(f"绘图时出错: {str(e)}")

def plot_transient_results(analysis):
    """
    绘制瞬态计算结果
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        # 创建一个包含4个子图的图形
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # 1. 残差历史
        ax1.semilogy(analysis.history['residuals'])
        ax1.set_title('残差历史')
        ax1.set_xlabel('时间步')
        ax1.set_ylabel('残差')
        ax1.grid(True)
        
        # 2. 最大空泡率随时间的变化
        ax2.plot(analysis.history['time'], analysis.history['max_void_fraction'])
        ax2.set_title('最大空泡率随时间的变化')
        ax2.set_xlabel('时间 (s)')
        ax2.set_ylabel('最大空泡率')
        ax2.grid(True)
        
        # 3. 平均压力随时间的变化
        ax3.plot(analysis.history['time'], analysis.history['avg_pressure'])
        ax3.set_title('平均压力随时间的变化')
        ax3.set_xlabel('时间 (s)')
        ax3.set_ylabel('压力 (Pa)')
        ax3.grid(True)
        
        # 4. 最大温度随时间的变化
        ax4.plot(analysis.history['time'], analysis.history['max_temperature'])
        ax4.set_title('最大温度随时间的变化')
        ax4.set_xlabel('时间 (s)')
        ax4.set_ylabel('温度 (K)')
        ax4.grid(True)
        
        plt.tight_layout()
        plt.show()
        
        # 创建动画展示场分布随时间的变化
        animate_field_evolution(analysis)
        
    except Exception as e:
        print(f"绘制瞬态结果时出错: {str(e)}")

def animate_field_evolution(analysis):
    """
    创建场分布随时间演化的动画
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        # 创建动画图形
        fig = plt.figure(figsize=(15, 5))
        
        # 创建三个子图：空泡率、温度和压力
        ax1 = fig.add_subplot(131, projection='3d')
        ax2 = fig.add_subplot(132, projection='3d')
        ax3 = fig.add_subplot(133, projection='3d')
        
        x = np.arange(analysis.n_axial)
        y = np.arange(analysis.n_channels)
        X, Y = np.meshgrid(x, y)
        
        # 确保有足够的历史数据
        n_frames = min(len(analysis.history['time']),
                      len(analysis.history['void_fraction_history']),
                      len(analysis.history['temperature_history']),
                      len(analysis.history['pressure_history']))
        
        if n_frames == 0:
            print("警告：没有可用的历史数据来创建动画")
            return
            
        def update(frame):
            # 清除当前帧
            ax1.clear()
            ax2.clear()
            ax3.clear()
            
            # 确保frame不会超出数组范围
            frame = min(frame, n_frames - 1)
            
            try:
                # 绘制空泡率分布
                surf1 = ax1.plot_surface(X, Y, 
                    analysis.history['void_fraction_history'][frame].T,
                    cmap=cm.viridis)
                ax1.set_title(f'空泡率分布\n时间: {analysis.history["time"][frame]:.3f} s')
                
                # 绘制温度分布
                surf2 = ax2.plot_surface(X, Y,
                    analysis.history['temperature_history'][frame].T,
                    cmap=cm.viridis)
                ax2.set_title('温度分布 (K)')
                
                # 绘制压力分布
                surf3 = ax3.plot_surface(X, Y,
                    analysis.history['pressure_history'][frame].T,
                    cmap=cm.viridis)
                ax3.set_title('压力分布 (Pa)')
                
                # 设置视角
                for ax in [ax1, ax2, ax3]:
                    ax.view_init(elev=20, azim=45)
                    ax.grid(True)
                    ax.set_xlabel('轴向位置')
                    ax.set_ylabel('子通道编号')
                    
            except IndexError as e:
                print(f"警告：动画帧 {frame} 超出数据范围")
                return
        
        # 创建动画
        anim = FuncAnimation(fig, update, 
                            frames=n_frames,
                            interval=200)
        
        # 保存动画
        try:
            anim.save(f'{analysis.results_folder}/evolution.mp4', 
                     writer='ffmpeg', fps=10)
            print(f"动画已保存到 {analysis.results_folder}/evolution.mp4")
        except Exception as e:
            print(f"保存动画失败: {str(e)}")
            print("尝试使用其他格式保存...")
            try:
                anim.save(f'{analysis.results_folder}/evolution.gif', 
                         writer='pillow', fps=10)
                print(f"动画已保存为GIF格式")
            except Exception as e:
                print(f"保存GIF也失败: {str(e)}")
        
        plt.tight_layout()
        plt.show()
        
    except Exception as e:
        print(f"创建动画时出错: {str(e)}")

def plot_channel_type_comparison(analysis):
    """
    绘制不同类型子通道的参数对比
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        # 创建图形
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # 获取各类型子通道的索引
        corner_channels = [j for j in range(analysis.n_channels) 
                         if analysis.channel_types[j+1] == 'corner']
        edge_channels = [j for j in range(analysis.n_channels) 
                        if analysis.channel_types[j+1] == 'edge']
        center_channels = [j for j in range(analysis.n_channels) 
                         if analysis.channel_types[j+1] == 'center']
        
        # 1. 空泡率分布对比
        ax1.plot(np.mean(analysis.void_fraction[:, corner_channels], axis=1), 
                label='角通道')
        ax1.plot(np.mean(analysis.void_fraction[:, edge_channels], axis=1), 
                label='边通道')
        ax1.plot(np.mean(analysis.void_fraction[:, center_channels], axis=1), 
                label='中心通道')
        ax1.set_title('空泡率轴向分布对比')
        ax1.set_xlabel('轴向位置')
        ax1.set_ylabel('平均空泡率')
        ax1.grid(True)
        ax1.legend()
        
        # 2. 温度分布对比
        ax2.plot(np.mean(analysis.T_l[:, corner_channels], axis=1), 
                label='角通道')
        ax2.plot(np.mean(analysis.T_l[:, edge_channels], axis=1), 
                label='边通道')
        ax2.plot(np.mean(analysis.T_l[:, center_channels], axis=1), 
                label='中心通道')
        ax2.set_title('液相温度轴向分布对比')
        ax2.set_xlabel('轴向位置')
        ax2.set_ylabel('平均温度 (K)')
        ax2.grid(True)
        ax2.legend()
        
        # 3. 液相速度分布对比
        ax3.plot(np.mean(analysis.liquid_velocity[:, corner_channels], axis=1), 
                label='���通道')
        ax3.plot(np.mean(analysis.liquid_velocity[:, edge_channels], axis=1), 
                label='边通道')
        ax3.plot(np.mean(analysis.liquid_velocity[:, center_channels], axis=1), 
                label='中心通道')
        ax3.set_title('液相速度轴向分布对比')
        ax3.set_xlabel('轴向位置')
        ax3.set_ylabel('平均速度 (m/s)')
        ax3.grid(True)
        ax3.legend()
        
        # 4. 压力分布对比
        ax4.plot(np.mean(analysis.pressure[:, corner_channels], axis=1), 
                label='角通道')
        ax4.plot(np.mean(analysis.pressure[:, edge_channels], axis=1), 
                label='边通道')
        ax4.plot(np.mean(analysis.pressure[:, center_channels], axis=1), 
                label='中心通道')
        ax4.set_title('压力轴向分布对比')
        ax4.set_xlabel('轴向位置')
        ax4.set_ylabel('平均压力 (Pa)')
        ax4.grid(True)
        ax4.legend()
        
        plt.tight_layout()
        plt.show()
        
        # 保存图形
        fig.savefig(f'{analysis.results_folder}/channel_type_comparison.png', 
                   dpi=300, bbox_inches='tight')
        print(f"对比图已保存到 {analysis.results_folder}/channel_type_comparison.png")
        
    except Exception as e:
        print(f"绘制���通道类型对比图时出错: {str(e)}")

def plot_convergence_history(analysis):
    """
    绘制收敛历史
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        # 创建图形
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # 1. 残差历史（对数坐标）
        ax1.semilogy(analysis.history['residuals'], 'b-', label='残差')
        ax1.set_title('残差收敛历史')
        ax1.set_xlabel('迭代次数')
        ax1.set_ylabel('残差')
        ax1.grid(True)
        ax1.legend()
        
        # 2. 迭代次数历史
        ax2.plot(analysis.history['time'], analysis.history['iterations'], 'r-', 
                label='迭代次数')
        ax2.set_title('每个时间步的迭代次数')
        ax2.set_xlabel('时间 (s)')
        ax2.set_ylabel('迭代次数')
        ax2.grid(True)
        ax2.legend()
        
        plt.tight_layout()
        plt.show()
        
        # 保存图形
        fig.savefig(f'{analysis.results_folder}/convergence_history.png', 
                   dpi=300, bbox_inches='tight')
        print(f"收敛历史图已保存到 {analysis.results_folder}/convergence_history.png")
        
    except Exception as e:
        print(f"绘制收敛历史时出错: {str(e)}")

def plot_heat_transfer_analysis(analysis):
    """
    绘制传热分析图
    
    参数:
    analysis: SubchannelAnalysis对象
    """
    try:
        # 创建图形
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
        
        # 1. 壁面温度分布
        im1 = ax1.contourf(analysis.T_c.T, cmap=cm.hot)
        ax1.set_title('壁面温度分布')
        ax1.set_xlabel('轴向位置')
        ax1.set_ylabel('子通道编号')
        fig.colorbar(im1, ax=ax1, label='温度 (K)')
        
        # 2. 热流密度分布
        im2 = ax2.contourf(analysis.q_wl.T, cmap=cm.hot)
        ax2.set_title('液相壁面热流密度分布')
        ax2.set_xlabel('轴向位置')
        ax2.set_ylabel('子通道编号')
        fig.colorbar(im2, ax=ax2, label='热流密度 (W/m²)')
        
        # 3. 沸腾曲线
        # 计算过热度
        delta_T = analysis.T_c - analysis.T_l
        # 计算热流密度
        q = analysis.q_wl
        
        ax3.scatter(delta_T.flatten(), q.flatten(), alpha=0.5, s=1)
        ax3.set_title('沸腾曲线')
        ax3.set_xlabel('壁面过热度 (K)')
        ax3.set_ylabel('热流密度 (W/m²)')
        ax3.set_yscale('log')
        ax3.grid(True)
        
        # 4. 换热系数分布
        h = q / delta_T
        h[np.isnan(h)] = 0
        h[np.isinf(h)] = 0
        
        im4 = ax4.contourf(h.T, cmap=cm.viridis)
        ax4.set_title('换热系数分布')
        ax4.set_xlabel('轴向位置')
        ax4.set_ylabel('子通道编号')
        fig.colorbar(im4, ax=ax4, label='换热系数 (W/m²·K)')
        
        plt.tight_layout()
        plt.show()
        
        # 保存图形
        fig.savefig(f'{analysis.results_folder}/heat_transfer_analysis.png', 
                   dpi=300, bbox_inches='tight')
        print(f"传热分析图已保存到 {analysis.results_folder}/heat_transfer_analysis.png")
        
    except Exception as e:
        print(f"绘制传热分析图时出错: {str(e)}")
